from __future__ import division
import torch
from torchvision import transforms
import cv2
import numpy as np
from matplotlib import pyplot as plt
from numpy import random


def scale_down(src_size, size):
    w, h = size
    sw, sh = src_size
    if sh < h:
        w, h = float(w * sh) / h, sh
    if sw < w:
        w, h = sw, float(h * sw) / w
    return int(w), int(h)


def fixed_crop(src, x0, y0, w, h, size=None):
    out = src[y0:y0 + h, x0:x0 + w]
    if size is not None and (w, h) != size:
        out = cv2.resize(out, (size[0], size[1]), interpolation=cv2.INTER_CUBIC)
    return out


def center_crop(src, size):
    h, w = src.shape[0:2]
    new_w, new_h = scale_down((w, h), size)

    x0 = int((w - new_w) / 2)
    y0 = int((h - new_h) / 2)

    out = fixed_crop(src, x0, y0, new_w, new_h, size)
    return out


class Compose(object):
    """Composes several augmentations together.
    Args:
        transforms (List[Transform]): list of transforms to compose.
    Example:
        >>> augmentations.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])
    """

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image):
        for t in self.transforms:
            image = t(image)
        return image

class RandomCrop(object):
    def __init__(self, tg_size):
        '''
        :param ori_size: (h,w)  eg (132,106)
        :param tg_size:  (h', w')  eg. (112,96)
        '''
        self.tg_size = tg_size
    def __call__(self, image):
        h,w = image.shape
        new_h, new_w = self.tg_size
        x0 = random.randint(0, w - new_w)
        y0 = random.randint(0, h - new_h)

        image = fixed_crop(image, x0, y0, new_w, new_h)
        return image

class RandomMirror(object):
    def __call__(self, image):
        if random.randint(2):
            image = image[:, ::-1]

        return image

class Normalize(object):
    def __call__(self, image):
        image = (image / 255.0 - 0.5) * 2
        return image

class FaceAug(object):
    def __init__(self, tg_size=(112, 96)):
        self.augment = Compose([
            RandomCrop(tg_size=tg_size),
            RandomMirror(),
            Normalize(),
        ])

    def __call__(self, image):
        return self.augment(image)